/*
 *	Copyright (C) 2008 CRIMERE
 *	Copyright (C) 2008 Jean-Marc Mercier
 *	
 *	This file is part of OTS (optimally transported schemes), an open-source library
 *	dedicated to scientific computing. http://code.google.com/p/optimally-transported-schemes/
 *
 *	CRIMERE makes no representations about the suitability of this
 *	software for any purpose. It is provided "as is" without express or
 *	implied warranty.
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *   You should have received a copy of the GNU General Public License
 *   along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

#if !defined _OS_tree_
#define _OS_tree_
#include <src/math/functor/GaussianFunctor.h>
#include <src/data/OS_Multi_Array.h>
#include <src/math/Stochastic/StochasticProcesses.h>
#include <src/math/distributions/distributions.h>

//! prototype class for stochastic trees.
/*! This class describes what we call a stochastic tree : 
	let \f$t \rightarrow mathrm{x}_t \in \Omega\f$ a stochastic variable, with density \f$\rho(t,x)\f$. 
	A stochastic tree corresponds to a sequence of sampling points and densities \f$\{S_i(t) \in \Omega,\rho_i(t) \in \R\}_{i=1,..,N}\f$
	to compute expectation of any function under the stochastic variable, that is for any function \f$\varphi\f$, the expectation
	is approxiamted accordingly to
    \f[
			\EE_{mathrm{x}_t}(\varphi) = \int_{\Omega}\varphi \rho(t,x)dx \sim \sum_{i=1,..,N} \rho_i(t) \varphi(S_i(t)) 
    \f]
*/

template <class stochastic_process>
struct OS_Tree_Base {
	typedef OS_Tree_Base													type;
	typedef stochastic_process												stochastic_process_type;
	typedef	typename stochastic_process_type::time_type						time_type;
	typedef	typename stochastic_process_type::underlying_type				underlying_type;
	typedef	typename stochastic_process_type::matrix_type					matrix_type;
	typedef	typename stochastic_process_type::base_value_type				base_value_type;

	enum :OS_size{Dim = stochastic_process::Dim};

	typedef typename OS_Multi_Array<Unbounded_Vector,
									Dim+1,
									base_value_type>::type
													underlying_vector_type;
	typedef typename OS_Multi_Array<Dirichlet_Vector,
									Dim+1,
									base_value_type>::type
													density_vector_type;

//! Ctor set the number of sampling points to zero.
	type() : steps_(0){};
	type(smart_ptr<stochastic_process_type> & process) : process_(process),steps_(0){};

//! prototype function to retrieve the set of sampling points \f$\{S_i(t) \in \Omega \}_{i=1,..,N}\f$. Input is understood as a time description.
	virtual underlying_vector_type underlying(const time_type & time) const {
		OS_DYNAMIC_CHECK(false, "to implement in derived class");  	
	};

//! prototype function to retrieve the set of densities \f$\{\rho_i(t) \in \R\}_{i=1,..,N}\f$. Input is understood as a time description.
	virtual density_vector_type densities(const time_type& time) const {
		OS_DYNAMIC_CHECK(false, "to implement in derived class");  	
	};
//! A default implementation to expectation computations.
	template<class function>
	OS_double expectation(const function &f, const time_type &t) const {
		return jmmath::sum(densities<time,density_vector_type>(t)^f(underlying<time,underlying_vector_type>(t)));
	};
//! accessor to the number of sampling points steps_.
	OS_size get_steps() {return steps_;};
//! a virtual function to set the number of sampling points steps_.
	virtual void set_steps(OS_size how_much) {steps_ = how_much;};
//! accessor to the process variable process_.
	smart_ptr<stochastic_process_type> get_process() {return process_;};

//! stores the numbers \f$N\f$ of sampling points.
	OS_size steps_;
//! stores a smart pointer over the process description.
	smart_ptr<stochastic_process_type> process_;
};


//! prototype class for stochastic trees based over geometric brownian processes.
/*! Not implemented yet. We implementd ony the one dimensional case specialization.
*/
template <class stochastic_process>
class OS_tree : public OS_Tree_Base<stochastic_process>
{
	OS_STATIC_CHECK(false); // to implement
};

//! class implementing a stochastic trees based over one dimensional geometric brownian processes.
/*! This class implement a stochastic tree (see OS_Tree) for one dimensional 
	geometric brownian motions.
	let \f$t \rightarrow mathrm{x}_t \in \R^+\f$ be a stochastic variable describing a one dimensional 
	geometric brownian motion with initial value \f$S_0\f$, rate \f$\mu\f$ and volatility \f$\sigma\f$.
	Let us denote the function \f$x = S(t,y)\f$ from \f$[0,1]\f$ into \f$\Omega = \RR+\f$ 
\f[
    S(t,y) = S_0 \exp \left( S_{\mathcal{N}}\left( \sigma^2 t,y \right)+\left(r-\frac{1}{2}\sigma^2\right) t \right)
\f]
	where \f$S_{\mathcal{N}}\left( t,y \right)\f$ denotes the inverse of the cumulative of the gaussian kernel
\f[
    \mathcal{N}(t,x) := \frac{1}{\sqrt{2 \pi t}} \exp\left( - \frac{x^2}{2t}\right)
\f]
	This tree implement the equiprobable sampling
    \f[
			\rho_i = \frac{1}{N} \text{ ; }S_i(t) = S(t,y_i) \text{ with }y_i=\frac{i+0.5}{N+1}
    \f]
	Property : we have the convergence rate
    \f[
			\EE_{mathrm{x}_t}(\varphi) = \sum_{i=1,..,N} \frac{1}{N}\varphi(S_i(t)) + \frac{1}{N}C(\varphi)
    \f]
	where \f$C(\varphi) = \sigma \sqrt{t} \|\varphi'\|_{L^{\infty}}\f$
*/
template <class geometric_process_type>
class OS_tree_first_order : public OS_Tree_Base<geometric_process_type>
	{
      public:
		  typedef OS_Tree_Base				Base;
		  typedef OS_tree_first_order		type;

		  type() {};
		  type(boost::shared_ptr<stochastic_process_type>& process) : Base(process)
		 { };
//! set the number of sampling points steps_. precompute the set inverse_cumulative_ and normal_densities_.
		virtual void set_steps(OS_size how_much) {
			steps_ = how_much;
			
			inverse_cumulative_.resize(steps_);
			normal_densities_.resize(steps_);
			inverse_cumulative_ = InverseCumulativeFunctor<underlying_vector_type>(0.,1.)(steps_);
			normal_densities_   = OS_normal_distribution<OS_double>(0.,1.).operator ()<density_vector_type,underlying_vector_type>(inverse_cumulative_);
		};
//! returns the set 
//! \{S_i(t) = S(t,y_i)\}_{i=1,..,N}
        underlying_vector_type underlying(OS_double time) const {

			underlying_type x0 = process_->stochastic_process_type::OS_Base::initial_conditions();
			underlying_type drift = process_->stochastic_process_type::OS_Base::drift(time,x0); // this is (r_1/2 \sigma ^2)
			matrix_type volat = process_->stochastic_process_type::OS_Base::volatility(time,x0);
			underlying_vector_type return_(inverse_cumulative_);

			OS_Array_multiply(return_,volat*std::sqrt(time));
			OS_Array_plus(return_,drift*time);
			OS_Array_exp(return_);
			OS_Array_multiply(return_,x0);
			return return_;
        };
//! compute the densities
//! \f$\{\rho_i(t) = \rho(t,S_i)\}_{i=1,..,N}\f$ according to the formula
//! \f[
//!    \rho(t,S_i)  :=  \frac{\mathcal{N}(\sigma^2 t, S_{\mathcal{N}}\left( \sigma^2 t,y_i \right))}{S_i}
//! \f]
        density_vector_type densities(OS_double time) const {
			 OS_STATIC_CHECK(underlying_vector_type::Dim == 1);
			 OS_STATIC_CHECK(density_vector_type::Dim == 1);
			double volat = process_->volatility(time,process_->initial_conditions());
			underlying_vector_type return_(normal_densities_);
			OS_Array_multiply(return_,1./(volat*std::sqrt(time)) );
			return return_^(1./underlying(time));
        };
//! implementation for expectation computations.
		template<class function>
		OS_double expectation(const function &f, OS_double t) const {
			return jmmath::moment1(f(underlying<time,underlying_vector_type>(t)));
		};
//! storing \f$\{S_{\mathcal{N}}(1,y_{i})\}_i\f$ for optimizations.
		underlying_vector_type inverse_cumulative_; 
//! storing \f$\{\mathcal{N}(1,S_{\mathcal{N}}(1,y_{i}))\}_i\f$ for optimizations.
		density_vector_type normal_densities_;		
   };

//! class implementing a stochastic trees based over one dimensional geometric brownian processes.
/*! This class is a different implementation a stochastic tree (see OS_Tree) for one dimensional 
	geometric brownian motions.
	let \f$t \rightarrow mathrm{x}_t \in \R^+\f$ be a stochastic variable describing a one dimensional 
	geometric brownian motion with initial value \f$S_0\f$, rate \f$\mu\f$ and volatility \f$\sigma\f$.
	Let us denote the gaussian kernel
\f[
    \mathcal{N}(t,x) := \frac{1}{\sqrt{2 \pi t}} \exp\left( - \frac{x^2}{2t}\right)
\f]
	denotes \f$S_{\mathcal{N}}^{-1}(t,x)\f$ from \f$\RR\f$ into \f$[0,1]\f$ the cumulative of the gaussian kernel \f$\mathcal{N}\f$,
	and \f$S_{\mathcal{N}}\left( t,y \right)\f$ from \f$[0,1]\f$ into \f$\RR\f$ the inverse of the cumulative of the gaussian kernel.
	This tree implements an equiprobable sampling \f$(\rho_i(t) = \frac{1}{N})\f$ with
\f[
   S_i = N S_0 \exp\left( rt \right) \left( S_{\mathcal{N}}^{-1}\left(z_{i,i+1}-\sigma^2 t\right) - S_{\mathcal{N}}^{-1}\left(z_{i,i-1}-\sigma^2 t\right)\right)
\f]
where \f$z_{i,i-1} = S_{\mathcal{N}}\left( \sigma^2 t,\frac{i}{N} \right)\f$.
Property : expectation converges at a \f$\frac{1}{N^2}\f$ rate for piecewise smooth functions
\f[
	\EE_{mathrm{x}_t}(\varphi) = \sum_{i=1,..,N} \frac{1}{N}\varphi(S_i(t)) + \frac{1}{N^2}C(\varphi)
\f]
where \f$C(\varphi) = \sigma \sqrt{t} \|\varphi'\|_{L^{\infty}}\f$.
*/
	template <>
	class OS_tree<geometric_process<>> : public OS_tree_first_order<geometric_process<>>
	{
      public:
		  typedef OS_tree_first_order Base;
		  typedef OS_tree type;
		  typedef Base::underlying_vector_type underlying_vector_type;
		  typedef Base::density_vector_type density_vector_type;
		  type(){};
		  type(boost::shared_ptr<stochastic_process_type>& process) : Base(process)
		 {};
//! set the number of sampling points steps_. 
//! precompute the set inverse_cumulative_ and normal_densities_ according to
//! \f$\{S_i=S_{\mathcal{N}}(1,\frac{i}{N})\}_i\f$,\f$\{\rho_i=\mathcal{N}(S_i)\}_i\f$ for optimizations.
		virtual void set_steps(OS_size how_much) {
			steps_ = how_much;
			inverse_cumulative_.resize(steps_);
			typedef boost::random::UniformSegmentGenerator<OS_double> uniform_generator_type;
			uniform_generator_type uniform_(steps_,1);
			density_vector_type return_(steps_);
			uniform_.fill(inverse_cumulative_);
			inverse_cumulative_[steps_-1] -= 0.000001;

			inverse_cumulative_ = OS_inverse_cumulative_normal_distribution<OS_double>(0.,1.).operator()<underlying_vector_type,underlying_vector_type>(inverse_cumulative_);
			normal_densities_   = OS_normal_distribution<OS_double>(0.,1.).operator ()<density_vector_type,underlying_vector_type>(inverse_cumulative_);
			normal_densities_[steps_-1] = 0.;
		};

 //! returns the set 
//! \{S_i(t) = S(t,y_i)\}_{i=1,..,N}
       underlying_vector_type underlying(OS_double time) const {
			OS_double x0 = process_->stochastic_process_type::OS_Base::initial_conditions();
			double rate = process_->stochastic_process_type::OS_Base::rate(time,x0); // this is r
			double volat = process_->stochastic_process_type::OS_Base::volatility(time,x0);

			density_vector_type return_(inverse_cumulative_);
			OS_Array_multiply(return_,volat*std::sqrt(time));
			OS_Array_minus(return_,volat*volat*time);

			return_ = OS_cumulative_normal_distribution<double>(0.,volat*std::sqrt(time)).operator ()<density_vector_type,density_vector_type>(return_);
			return_ = (return_ - (return_ >> 1)) * return_.size();

			OS_Array_multiply(return_,x0*std::exp(rate * time));
			return return_;
        };
   };

//! class implementing a stochastic trees based imbedding a cauchy problem to compute the densities and sampling set.
/*! This class is a wrapper for stochastic trees. It may be used together with processes for which the sampling points and densities
necessitate a PDE solver to be computed.
Note : an explicit tree (class tree) is needed to settle the initial conditions of the Cauchy problem.
*/
	template <class stochastic_process, class cauchy, class initial_condition>
	class OS_transport_map_scheme : public OS_Tree_Base<stochastic_process>
	{
      public:
		typedef initial_condition initial_condition_type;
		typedef stochastic_process stochastic_process_type;
		typedef cauchy cauchy_type;

		typedef OS_Tree_Base Base;

		  OS_transport_map_scheme(boost::shared_ptr<stochastic_process_type>& process) : 
			Base(process)
			, init_(process_)
			, cauchy_(process_)
		 {};
//! The cauchy problem is solved from time/10 to time. The initial conditions are given by an explicit tree
//!	time/10.
	virtual underlying_vector_type underlying(const time_type & time) const {
				OS_double time_begin = time/10.;
				OS_double time_end = time;
				cauchy_.reset_time();
				cauchy_.insert_time(time_begin);
				cauchy_.insert_time(time_end);

				cauchy_.set_initial_conditions(init_(time_begin));
			
			for (;time_begin < time_end; time_begin = ++ cauchy_ ) {
//				std::cout << " Time =  " << time_begin << "   ";
//				std::cout << " moyenne =  " << jmmath::moment1(cauchy_.get_solution()) << "    ";
//				std::cout << " moyenne theo =  " << jmmath::moment1(tree_.underlying(time_begin)) << std::endl;
			}
//			std::cout << " moyenne theo =  " << jmmath::moment1(tree_.underlying(time_end)) << std::endl;
			return cauchy_.get_solution();
		};
        density_vector_type densities(OS_double time) const {
			OS_DYNAMIC_CHECK(false,"to check if coherent / to implement");
			return density_vector_type();
        };

      protected:
		initial_condition_type init_;
		mutable cauchy_type cauchy_;
    };

//! prototype class for stochastic trees based over Cox Ingelsonn Ross processes.
/*! prototype class for stochastic trees based over Cox Ingelsonn Ross processes.
*/
template <>
class OS_tree<CIR_process<>> : public OS_Tree_Base<CIR_process<>>
{
      public:
		typedef OS_Tree_Base	Base;
		typedef OS_tree			type;
		type(){};
		type(boost::shared_ptr<stochastic_process_type>& process) : Base(process)
		{};
//! returns the set 
//! \{S_i(t) = S(t,y_i)\}_{i=1,..,N}
		virtual underlying_vector_type underlying(const time_type & time) const {
			OS_double kappa = process_->kappa_;
			OS_double x0	= process_->S0_;
			OS_double sigma	= process_->sigma_;
			OS_double theta = process_->theta_;
			OS_double delta = process_->delta_;

			OS_double mu = (4.*kappa)/(sigma*sigma*(1.-exp(-kappa*time)));
			OS_double lambda = mu*exp(-kappa*time)*x0;

			OS_inverse_cumulative_non_central_chi_squared_distribution<OS_double> chi(delta,lambda);
			OS_cumulative_non_central_chi_squared_distribution<OS_double> chi_2(delta+2,lambda);
			OS_cumulative_non_central_chi_squared_distribution<OS_double> chi_4(delta+4,lambda);

			typedef boost::random::UniformSegmentGenerator<OS_double> uniform_generator_type;
			uniform_generator_type uniform_(steps_,1);
			density_vector_type temp(steps_);
			uniform_.fill(temp);
			temp[steps_-1] -= 0.000001;
			temp = chi.operator ()<density_vector_type,density_vector_type>(temp); // these are S_{\delta}^{\mu}(i/N)

			density_vector_type return_ = (chi_2.operator ()<density_vector_type,underlying_vector_type>(temp))*(delta)
				+(chi_4.operator ()<density_vector_type,underlying_vector_type>(temp))*lambda;
			return_ = (return_ - (return_ >> 1)) * (return_.size()/mu);
			
//			double test1 = jmmath::sum(return_)/return_.size();
//			double test2 = (delta+lambda)/mu; // should be close.

			return return_;
			
/* first order tree
			typedef CIR_process_transport_map<stochastic_process_type> CIR_transport_map_type;
			struct CIR_map_type : public template_function_utility<CIR_transport_map_type>
			{
				typedef template_function_utility Base;
				CIR_map_type(const smart_ptr<stochastic_process_type> & process, OS_double time) : Base(process,time){};
			};
			CIR_map_type CIR_process_transport_map_(process_,time);

			typedef boost::random::UniformEquallySpacedGenerator<OS_double> uniform_generator_type;
			uniform_generator_type uniform_(steps_);
			density_vector_type temp(steps_);
			uniform_.fill(temp);
			return CIR_process_transport_map_.operator()<underlying_vector_type,density_vector_type>(temp); 
*/
        };
};


//! A vector of  one dimensional stochastic trees prototype.
/*! This class behaves mainly as a vector of stochastic tree.
*/
template <class tree_type>
class OS_Vector_tree : public std::vector<tree_type>
{
      public:
		  typedef OS_Vector<typename tree_type::underlying_type>	underlying_type;
		  typedef std::vector<underlying_type>						underlying_vector_type;
		  typedef std::vector<tree_type> Base;
		  typedef OS_Vector<OS_size> size_vector_type;

		  OS_Vector_tree(OS_size n) : Base(n){};
		  OS_Vector_tree() {};


//! calls each trees and set its number of steps to N. 
		virtual void set_steps(OS_size how_much) {
			OS_size size_(size());
			size_list_.resize(size_);
			size_list_.assign(size_,how_much);
			for (OS_size i = 0;i<size_;++i) {
				(*this)[i].set_steps( size_list_[i] );
			};
		};

		underlying_vector_type underlying(double time) const {
			OS_size size_(size());
			underlying_vector_type result(size_);
			for (OS_size i = 0;i<size_;++i) {
				result[i] = (*this)[i].underlying(time);
			};
			return result;
		};
		
//! stores the number of sampling point of each tree.
		std::vector<OS_size>		size_list_;
};
template <class tree_type>
class OS_Cartesian_tree : public OS_Vector_tree<tree_type>
{
      public:
		  typedef OS_Vector_tree Base;

		  OS_Cartesian_tree(OS_size n) : Base(n){};
		  OS_Cartesian_tree() {};


		virtual void set_steps(OS_size how_much) {
			OS_size size_(size());
			size_list_.resize(size_);
			OS_double log_size = std::pow(how_much,1./(double) size_ );
			OS_size size_dim = (1 > log_size ? 1 : log_size);
			size_list_.assign(size_,size_dim);

			OS_size log_size_inf = 1.;
			for (int i = 0;i<size_list_.size();++i)
				log_size_inf *= size_list_[i];

			std::vector<OS_size>::iterator size_list_it = size_list_.begin();
			while (log_size_inf < how_much) {
				++(*size_list_it);
				log_size_inf = 1.;
				for (int i = 0;i<size_list_.size();++i)
					log_size_inf *= size_list_[i];
				++size_list_it;
			};

			for (OS_size i = 0;i<size_;++i) {
				(*this)[i].set_steps( size_list_[i] );
			};
		};
};

/*! helper. Defines a generator with a tree for MonteCarlo methods.
*/
template<class payoff_type, class tree_type>
struct tree_generator : public TrivialGenerator<typename tree_type::underlying_vector_type> 
{
	typedef TrivialGenerator Base;
	typedef tree_generator type;
	typedef typename tree_type::stochastic_process_type stochastic_process_type;
	type(){};
	type(const type & from) : Base(from) {};
	type(smart_ptr<stochastic_process_type> & process, smart_ptr<payoff_type> & payoff) : 
	tree_(process),payoff_(payoff)
	{};
	const type& operator = (const type & from) {
		Base::operator = (from);
		payoff_ = from.payoff_;
		tree_	= from.tree_;
		return *this;
	};
	void set_steps (unsigned long Samples_nb) {
		tree_.set_steps(Samples_nb);
		set_data(tree_.underlying(payoff_->get_maturity()));
	};
	smart_ptr<payoff_type>				payoff_;
	tree_type tree_;
};


#endif